package com.zengde.ai_sandbox.config;

import com.zengde.ai_sandbox.handler.EchoWebSocketHandler;
import com.zengde.ai_sandbox.handler.RawWebSocketHandler;
import org.springframework.boot.autoconfigure.condition.ConditionalOnWebApplication;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Profile;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.lang.NonNull;
import org.springframework.lang.Nullable;
import org.springframework.messaging.Message;
import org.springframework.messaging.MessageChannel;
import org.springframework.messaging.simp.config.ChannelRegistration;
import org.springframework.messaging.simp.config.MessageBrokerRegistry;
import org.springframework.messaging.simp.stomp.StompCommand;
import org.springframework.messaging.simp.stomp.StompHeaderAccessor;
import org.springframework.messaging.support.ChannelInterceptor;
import org.springframework.messaging.support.MessageHeaderAccessor;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.config.annotation.*;
import org.springframework.web.socket.server.HandshakeInterceptor;
import org.springframework.web.socket.server.standard.ServletServerContainerFactoryBean;

import java.util.Map;
import java.util.concurrent.TimeUnit;

@Configuration
@EnableWebSocketMessageBroker
@EnableWebSocket
@ConditionalOnWebApplication(type = ConditionalOnWebApplication.Type.SERVLET)
@Profile("!test")
public class WebSocketConfig implements WebSocketMessageBrokerConfigurer, WebSocketConfigurer {
    
    private final RawWebSocketHandler rawWebSocketHandler;
    private final EchoWebSocketHandler echoWebSocketHandler;
    
    public WebSocketConfig(RawWebSocketHandler rawWebSocketHandler, EchoWebSocketHandler echoWebSocketHandler) {
        this.rawWebSocketHandler = rawWebSocketHandler;
        this.echoWebSocketHandler = echoWebSocketHandler;
    }

    @Override
    public void configureMessageBroker(@NonNull MessageBrokerRegistry registry) {
        // 客户端订阅消息的前缀
        registry.enableSimpleBroker("/topic", "/queue");
        // 客户端发送消息的前缀
        registry.setApplicationDestinationPrefixes("/app");
        // 设置用户目的地前缀
        registry.setUserDestinationPrefix("/user");
    }

    @Override
    public void registerStompEndpoints(@NonNull StompEndpointRegistry registry) {
        // 添加一个握手拦截器，记录WebSocket连接信息
        HandshakeInterceptor loggingInterceptor = new HandshakeInterceptor() {
            @Override
            public boolean beforeHandshake(@NonNull ServerHttpRequest request, @NonNull ServerHttpResponse response, 
                                        @NonNull WebSocketHandler wsHandler, @NonNull Map<String, Object> attributes) {
                System.out.println("WebSocket连接握手开始: " + request.getURI());
                return true;
            }

            @Override
            public void afterHandshake(@NonNull ServerHttpRequest request, @NonNull ServerHttpResponse response, 
                                    @NonNull WebSocketHandler wsHandler, @Nullable Exception exception) {
                System.out.println("WebSocket连接握手完成: " + request.getURI());
            }
        };
        
        // 注册一个WebSocket端点，客户端通过这个端点连接到WebSocket服务器
        registry.addEndpoint("/ws")
                .setAllowedOriginPatterns("*")  // 允许所有域名访问
                .addInterceptors(loggingInterceptor);
        
        // 同时支持SockJS连接
        registry.addEndpoint("/ws")
                .setAllowedOriginPatterns("*")
                .withSockJS()  // 使用SockJS
                .setInterceptors(loggingInterceptor);
    }

    @Override
    public void registerWebSocketHandlers(@NonNull WebSocketHandlerRegistry registry) {
        System.out.println("注册WebSocket处理器...");
        
        // 注册纯WebSocket处理器
        registry.addHandler(rawWebSocketHandler, "/wsraw")
                .setAllowedOrigins("*");  // 允许所有来源访问
        
        // 添加SockJS支持
        registry.addHandler(rawWebSocketHandler, "/wsraw")
                .setAllowedOrigins("*")
                .withSockJS();
        
        // 注册Echo WebSocket处理器 (用于简单测试)
        registry.addHandler(echoWebSocketHandler, "/echo")
                .setAllowedOrigins("*");
        
        // Echo处理器也添加SockJS支持
        registry.addHandler(echoWebSocketHandler, "/echo")
                .setAllowedOrigins("*")
                .withSockJS();
        
        System.out.println("WebSocket处理器注册完成");
    }

    @Override
    public void configureWebSocketTransport(@NonNull WebSocketTransportRegistration registration) {
        registration.setMessageSizeLimit(128 * 1024); // 限制消息大小为128KB
        registration.setSendBufferSizeLimit(512 * 1024); // 发送缓冲区大小为512KB
        registration.setSendTimeLimit(20000); // 发送超时时间为20秒
    }

    @Override
    public void configureClientInboundChannel(@NonNull ChannelRegistration registration) {
        registration.interceptors(new ChannelInterceptor() {
            @Override
            public Message<?> preSend(@NonNull Message<?> message, @NonNull MessageChannel channel) {
                StompHeaderAccessor accessor = MessageHeaderAccessor.getAccessor(message, StompHeaderAccessor.class);
                if (accessor != null) {
                    StompCommand command = accessor.getCommand();
                    if (command != null) {
                        System.out.println("收到STOMP消息: " + command + ", sessionId: " + accessor.getSessionId());
                        if (StompCommand.SUBSCRIBE.equals(command)) {
                            System.out.println("订阅地址: " + accessor.getDestination() + ", subscriptionId: " + accessor.getSubscriptionId());
                        }
                        else if (StompCommand.SEND.equals(command)) {
                            System.out.println("发送地址: " + accessor.getDestination());
                        }
                    }
                }
                return message;
            }
        });
    }

    @Override
    public void configureClientOutboundChannel(@NonNull ChannelRegistration registration) {
        registration.interceptors(new ChannelInterceptor() {
            @Override
            public Message<?> preSend(@NonNull Message<?> message, @NonNull MessageChannel channel) {
                StompHeaderAccessor accessor = MessageHeaderAccessor.getAccessor(message, StompHeaderAccessor.class);
                if (accessor != null) {
                    StompCommand command = accessor.getCommand();
                    if (command != null) {
                        System.out.println("发送STOMP消息: " + command + ", sessionId: " + accessor.getSessionId());
                        if (StompCommand.MESSAGE.equals(command)) {
                            System.out.println("消息地址: " + accessor.getDestination() + ", subscriptionId: " + accessor.getSubscriptionId());
                        }
                    }
                }
                return message;
            }
        });
    }

    @Bean
    // 添加条件注解，避免在测试环境中初始化WebSocket容器时出错
    @ConditionalOnWebApplication
    public ServletServerContainerFactoryBean createWebSocketContainer() {
        ServletServerContainerFactoryBean container = new ServletServerContainerFactoryBean();
        // 增加缓冲区大小和超时时间以提高性能
        container.setMaxTextMessageBufferSize(32768); // 32KB
        container.setMaxBinaryMessageBufferSize(32768); // 32KB
        container.setMaxSessionIdleTimeout(TimeUnit.HOURS.toMillis(1)); // 1 hour
        container.setAsyncSendTimeout(TimeUnit.SECONDS.toMillis(30)); // 30 seconds        // Set the maximum time that a WebSocket session can remain idle
        return container;
    }
}